tjbal_unbalance = function(data, outcome, baseline_year, 
                           estimator_type, vce){
  
  force_balance = function(data, outcome, baseline_year){
    
    data = data 
    data$outcome = outcome 
    
    data <- data %>% 
      group_by(DISTID) %>% 
      filter(year >= baseline_year) %>% 
      mutate(na_ever = mean(outcome)) %>% 
      filter(is.na(na_ever) == F) %>% 
      as.data.frame()
    
    return(data)
  }
  
  data = force_balance(data, outcome, baseline_year)
  
  m = tjbal(data, Y = 'outcome', 
            D = "treat", 
            index = c("DISTID", "year"), 
            demean = F, estimator = estimator_type, 
            vce = vce)
  
  return(m)
}

anqar_panel = readRDS('data/anqar_panel_2023.rds')
df <- readRDS('data/full_panel_df.rds') %>% 
  filter(cohort == "2011" | cohort == "2012" | cohort == '2013' | cohort == '10000') %>% 
  filter(year >= 2008) %>% 
  filter(year <= 2014) %>% 
  mutate(
    treat = ifelse(treated_ever==1, as.numeric(year>=cohort, 1,0), 0)
  ) %>% 
  select('DISTID', 'year', 'cohort', 'treat', 'treated_ever'
  ) 
df = merge(df, anqar_panel, by = c('DISTID', 'year'), 
            all.x = TRUE 
)

dispute_court_plot = make_tj_ct(tjbal_unbalance(df, df$dispute_court, 2009, 
                           'mean', 'none'), 
                           'Years until Taliban Courts', 'State Courts', 
                           '') + 
  theme_tufte() + 
  theme(legend.position = 'bottom')
ggsave("fig-out/dispute_court_plot.pdf")


TLB_return_plot = make_tj_ct(tjbal_unbalance(df, df$return_good, 2009, 
                                                'mean', 'none'), 
                                'Years until Taliban Courts', 'Taliban Approval', 
                                '') + 
  theme_tufte() + 
  theme(legend.position = 'bottom')
ggsave("fig-out/TLB_return_plot.pdf")


gov_indx_plot = make_tj_ct(tjbal_unbalance(df, df$gov_index2, 2008, 
                                             'mean', 'none'), 
                           'Years until Taliban Courts', 'Gov. Index', 
                           '') + 
  theme_tufte() + 
  theme(legend.position = 'bottom')
ggsave("fig-out/gov_indx_plot.pdf")


gov_cntr_plot = make_tj_ct(tjbal_unbalance(df, df$gov_influence, 2008, 
                                           'mean', 'none'), 
                           'Years until Taliban Courts', 'Gov. Control', 
                           '') + 
  theme_tufte() + 
  theme(legend.position = 'bottom')
ggsave("fig-out/gov_cntr_plot.pdf")

